# Copyright (c) 2023 Presto Labs Pte. Ltd.
# Author: jhkim

import datetime
import pandas
import os
import re
import json
import traceback
import tempfile

from absl import flags

from cc.appcoin2.strategy.sim_backtest import SimBacktest
from cc.appcoin2.strategy.run_sim_util import create_temp_file_and_dump_json
import coin2.app.driver_pb2 as app_pb2
import coin2.strategy.linear_model.lm_config_pb2 as lm_pb2
from python.coin_deploy.lm_agg2 import util as lmutil
from python.coin_deploy.lm_backtest.lm_sim_param_override import rmm_search_params


rmm_keys = [rmm_search_param[0] for rmm_search_param in rmm_search_params]

def get_item_recursive(obj, key):
  ret = None
  if isinstance(obj, list):
    for v in obj:
      ret = ret or get_item_recursive(v, key)
  elif isinstance(obj, dict):
    if key in obj:
      ret = ret or obj[key]
    for k, v in obj.items():
      if k == key:
        ret = ret or v
      else:
        ret = ret or get_item_recursive(v, key)
  return ret


def replace_item(obj, key, replace_value):
  if isinstance(obj, list):
    for v in obj:
      replace_item(v, key, replace_value)
  elif isinstance(obj, dict):
    if key in obj:
      obj[key] = replace_value
    for k, v in obj.items():
      replace_item(v, key, replace_value)


def pick_dict_keys(base_dict, pick_keys):
  pick_dict = {key: None for key in pick_keys}
  split_dict = {}
  for key, value in base_dict.items():
    if key in pick_keys:
      pick_dict[key] = value
    else:
      split_dict[key] = value
  return pick_dict, split_dict


class RmmBacktest(SimBacktest):
  def __init__(
        self,
        tmp_dirname,
        analyzer_feed_machine,
        config_path,
        *,
        sigcache=None,
        sim_feed_machine=None,
        **sim_args):
    self.pb_path_override = None
    self.rmm_sim_args, base_sim_args = pick_dict_keys(sim_args, rmm_keys)
    self.tmp_dirname = tmp_dirname

    sim_feed_machine = "feed-05.ap-northeast-1.aws"

    if flags.FLAGS.debug:
      flags.FLAGS.fatsim = 0

    SimBacktest.__init__(
        self,
        tmp_dirname,
        analyzer_feed_machine,
        config_path=config_path,
        sim_feed_machine=sim_feed_machine,
        **base_sim_args)

    os.makedirs(os.path.join(self.sim_result_dir, "log"), exist_ok=True)
    if len(self.trading_intervals) > 0:
      start_date = self.trading_intervals[0][0].strftime('%Y%m%d')
      end_date = (
          self.trading_intervals[-1][-1] -
          datetime.timedelta(minutes=1)).strftime('%Y%m%d')
      self.strat_prefix = f"{start_date}-{end_date}_" + self.strat_prefix

    if flags.FLAGS.note is not None and len(flags.FLAGS.note) > 0:
      self.strat_prefix = f"{flags.FLAGS.note}_{self.strat_prefix}"

    self.strat_prefix_base = self.strat_prefix

    for key, short in rmm_search_params:
      val = self.rmm_sim_args[key]
      if val is not None:
        if "nthprod" == key:
          self.strat_prefix_base += f"_{short}_{val}"
          continue
        if key.endswith("ns"):
          valms = int(val * 1e-6)
          self.strat_prefix += f"_{short}_{valms}ms"
        elif key.find("regex") >= 0:
          for escape in [".", "*", "|", "(", ")", "$", "^", "&"]:
            val = val.replace(escape, "")
          self.strat_prefix_base += f"_{short}_{val}"
        else:
          self.strat_prefix += f"_{short}_{val}"

    self.csv_suffix = '*.csv'
    self.initial_position = sim_args.get("initial_position", 0)
    self.strat_prefix += "_"

    strat_config = json.load(open(os.path.join(os.environ.get("COIN_REPO"), flags.FLAGS.config_path)))
    driver = json.load(open(os.path.join(os.environ.get("COIN_REPO"), flags.FLAGS.driver_path)))

    del(driver['feed']['common']['raw_parser'])
    driver['feed']['common']['type'] = "FEED_SYSTEM_TYPE_ARCHIVE"
    driver['feed']['common']['archive'] = {}
    driver['feed']['common']['archive']['machine'] = sim_feed_machine
    driver['feed']['common']['archive']['archive_type'] = 'FASTFEED'
    for key, val in driver['feed']['exchanges'].items():
      if "websocket" in val:
        del val['websocket']
      val['type'] = "FEED_SYSTEM_TYPE_ARCHIVE"
      if key == "Futures.Prex.v1.realtime":
        symbol = val['products']['norms'][0]
    driver['order']['common'] = {}
    driver['order']['common']['sim_config2'] = {}
    driver['order']['common']['sim_config2']['accounts'] = []
    driver['order']['common']['sim_config2']['accounts'].append(
            {"relative_norm": symbol})
    if 'strategy_health_reporter_config' in driver['strategy']:
      del driver['strategy']['strategy_health_reporter_config']
    driver['order']['exchanges']['Prex']['type'] = 'ORDER_SYSTEM_TYPE_SIM_OG'
    self.driver_filename = create_temp_file_and_dump_json(
        dir=self.tmp_dirname, prefix="driver.", suffix=".json", data=driver)
    self.strat_config_filename = create_temp_file_and_dump_json(
        dir=self.tmp_dirname, prefix="strat_config.", suffix=".json", data=strat_config)
    self.key_filename = os.path.join(os.environ.get("COIN_REPO"), flags.FLAGS.key_path)
    self.sim_binary = "../coin_binary/coin2-bazel-bin-3.10/cc/appcoin2/strategy/rmm/rmm_runner"
    self.sim_account_id = 0

  def gen_cmds_prerun(self):
    return []

  def gen_cmds_kline(self):
    return []

  def cmd_run_on_date(self, trading_interval, i):
    tslist = pandas.DatetimeIndex(trading_interval).astype(int).tolist()
    from_ts = tslist[0]
    to_ts = tslist[1]
    pipecmd = ""

    if flags.FLAGS.valgrind or flags.FLAGS.debug == 2:
      sim_binary = f"valgrind --tool=callgrind {self.sim_binary}"
    else:
      sim_binary = f"{self.sim_binary}"
    from_dt_str = trading_interval[0].strftime("%Y%m%dT%H%M%S")
    from_ts = int(trading_interval[0].timestamp() * 1e9)
    log_filepath = None

    if not flags.FLAGS.debug:
      if len(self.strat_prefix) > 200:
        import hashlib
        logpref = hashlib.md5(self.strat_prefix.encode('utf-8')).hexdigest()
        log_filepath = f"{self.sim_result_dir}/log/{logpref}_{from_dt_str}.txt"
      else:
        log_filepath = f"{self.sim_result_dir}/log/{self.strat_prefix}_{from_dt_str}.txt"
      pipecmd = f"|& (head -n 500; tail -n 500) &>> {log_filepath}"

    start = pandas.to_datetime(trading_interval[0])
    end = pandas.to_datetime(trading_interval[1]) - pandas.Timedelta('1s')
    driver = json.load(open(self.driver_filename))
    driver['feed']['common']['archive']['interval'] = {
        'start': start.strftime('%Y%m%dT%H%M%S'),
        'end': end.strftime('%Y%m%dT%H%M%S'),
    }
    if trading_interval[0] < datetime.datetime(2022, 11, 11):
      driver['feed']['exchanges']["Futures.Prex.v1.usdtusd_index"]['mea'] = "Spot.Binance.v1"
      driver['feed']['exchanges']["Futures.Prex.v1.usdtusd_index"]['products']['norms'] = ['BUSD-USDT']
      driver['feed']['exchanges']["Futures.Prex.v1.usdtusd_index"]['recipe'] = "l1_realtime"
    driver['order']['common']['sim_config2']['accounts'][0]['start_timestamp'] = from_ts

    strat_config = json.load(open(self.strat_config_filename))
    for k, v in self.rmm_sim_args.items():
      if v is not None:
        strat_config['product_configs'][0][k] = v
    self.driver_filename_tmp = create_temp_file_and_dump_json(
        dir=self.tmp_dirname, prefix="driver.", suffix=".json", data=driver)
    self.strat_config_filename_tmp = create_temp_file_and_dump_json(
        dir=self.tmp_dirname, prefix="strat_config.", suffix=".json", data=strat_config)
    self.sim_info_filename = self.strat_config_filename_tmp.replace(
        "/strat.", "/sim_info.")
    self.sparams_filename = self.strat_config_filename_tmp.replace(
        "/strat.", "/sparams.")
    if flags.FLAGS.debug == 3:
      cmd_sim = f"""cat {self.sparams_filename}; \
          cat {self.strat_config_filename_tmp}; cat {self.driver_filename_tmp}"""
    else:
      cmd_sim = [f"export GLOG_log_dir={self.tmp_dirname}/;"]
      cmd_sim += [sim_binary]
      flag_dict = { \
        "driver": self.driver_filename_tmp,
        "config": self.strat_config_filename_tmp,
        "key_filepath": self.key_filename,
        "sim_result_filepath": f"{self.tmp_dirname}/{from_dt_str}.csv",
        "sim_start_ts": from_ts,
        "is_sim": 1}
      cmd_sim += [f"--{key}={val}" for key, val in flag_dict.items()]
      cmd_sim += [pipecmd]
      cmd_sim = " ".join(cmd_sim)
      if i == 0:
        print(f"first cmd sim:\n{cmd_sim}")
    if log_filepath:
      with open(log_filepath, 'w') as fwrt:
        fwrt.write(cmd_sim)
        fwrt.write("\n")
        fwrt.close()
    return cmd_sim

  def cmd_analyze_and_plot(self, compress=False, cpus=None):
    cpus_use = len(self.trading_intervals) if not cpus else cpus
    cpus_use = min(cpus_use, 6)
    if len(self.trading_intervals) == 1:
      trading_interval = self.trading_intervals[0]
      fromdt, todt = trading_interval[0], trading_interval[1]
      tss = pandas.DatetimeIndex([fromdt, todt]).astype(int)
    sim_info_str = ""
    if self.sim_info_filename:
      sim_info_str = f"--sim_info {self.sim_info_filename}"
    cmd = f"""./pyrunner \
python/coin/strategy/rmm/sim_stat_fromcsv.py \
--csv_path "{self.csv_path()}" \
--sim_result_prefix={self.strat_prefix} \
--strat_prefix={self.strat_prefix_base} \
--trading_date={self.trading_date} \
--time_range={self.time_range} \
--sim_result_dir={self.sim_result_dir} \
--cpu={cpus_use} \
{sim_info_str}"""
    return cmd, cpus_use

  def csv_path(self):
    return (
        self.pb_path_override or
        f"{self.tmp_dirname}/{self.csv_suffix}")
